import logging
from datetime import datetime
from tqdm import tqdm
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from pinecone import Pinecone, ServerlessSpec

# ---------- 配置日志 ----------
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# ---------- 初始化 Pinecone ----------
pc = Pinecone(api_key="pcsk_4RZEQz_7gy8B1tBbXWFw22UX6swnKLafVduhYCrTx4eULaxrpgsWEEAnt96Y7hXstdgFqi")
index_name = "mnist-index"

# ---------- 加载 MNIST 数据集 ----------
digits = load_digits()
X, y = digits.data, digits.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# ---------- 删除已有索引（如果存在） ----------
if index_name in pc.list_indexes().names():
    pc.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已存在，已删除。")

# ---------- 创建新索引 ----------
pc.create_index(
    name=index_name,
    dimension=64,
    metric="euclidean",
    spec=ServerlessSpec(cloud="aws", region="us-east-1")
)
logging.info(f"索引 '{index_name}' 创建成功。")
index = pc.Index(index_name)

# ---------- 上传训练数据 ----------
logging.info("开始上传数据到 Pinecone...")
vectors = [(str(i), X_train[i].tolist(), {"label": int(y_train[i])}) for i in range(len(X_train))]
for i in tqdm(range(0, len(vectors), 100), desc="上传数据到Pinecone"):
    index.upsert(vectors[i:i+100])

logging.info("成功上传了 %d 条数据到 Pinecone。", len(X_train))

# ---------- 测试 K=11 的准确率 ----------
k = 11
correct = 0
logging.info("测试 k=%d 的准确率...", k)

for i in tqdm(range(len(X_test)), desc="测试k=11的准确率"):
    query_vector = X_test[i].tolist()
    result = index.query(vector=query_vector, top_k=k, include_metadata=True)

    if not result.matches:
        continue  # 如果没返回结果，跳过

    labels = [int(match.metadata["label"]) for match in result.matches]
    pred = max(set(labels), key=labels.count)
    if pred == y_test[i]:
        correct += 1

accuracy = correct / len(y_test)
logging.info("K=%d 时的准确率为：%.4f", k, accuracy)